import cv2
import torch
from pathlib import Path
import numpy as np
from glob import glob
from tqdm import tqdm
from models.assembly.segmentation_table import Segmentation_Model


def save_tensor(tensor, i, im):
    im = cv2.resize(im, (1280, 800))
    np_array = tensor[0].numpy().transpose(1, 2, 0)
    np_array = cv2.resize(np_array, (1280, 800))
    im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    im[np_array > 0.3] = 255
    # new_im = im * np_array
    cv2.imwrite(str(i)+'.jpg', im)
    # cv2.imwrite(str(i)+'.jpg', np.vstack((np_array * 254, im)))


def main():
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    model = Segmentation_Model().to(device)
    model.load_state_dict(torch.load('../weights_trained/99.pth', map_location=torch.device('cpu')))
    model.eval()
    fd = r'F:\datasets\ocr\internal\sg'
    images = [str(i) for i in Path(fd).glob('*.jp*')]
    width, height = 1280, 800
    with torch.no_grad():
        for i, image in tqdm(enumerate(images)):
            im1 = cv2.imread(image)
            im = cv2.resize(im1, (width, height))/255/ 255.0
            im = im[:, :, ::-1].transpose(2, 0, 1)
            im = np.ascontiguousarray(im)
            im = torch.from_numpy(im).float()
            if im.ndimension() == 3:
                im = im.unsqueeze(0)
            pred = model(im)
            save_tensor(pred, i, im1)


if __name__ == '__main__':
    main()
